'''
Author: your name
Date: 1970-01-01 08:00:00
LastEditTime: 2020-12-10 10:24:22
LastEditors: Please set LastEditors
Description: In User Settings Edit
FilePath: /Pointnet_Pointnet2/models/center_loss.py
'''

import torch
import torch.nn as nn
import torch.nn.functional as F

class get_loss(nn.Module):
    """Center loss.
    
    Reference:
    Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016.
    
    Args:
        cls_num (int): number of classes.
        feat_dim (int): feature dimension.
    """
    def __init__(self, args, use_gpu=True):
        super(get_loss, self).__init__()
        self.cls_num = args.num_class
        self.feat_dim = args.feat_dim
        self.use_gpu = use_gpu

        if self.use_gpu:
            self.centers = nn.Parameter(torch.randn(self.cls_num, self.feat_dim).cuda())
        else:
            self.centers = nn.Parameter(torch.randn(self.cls_num, self.feat_dim))

    def forward(self, x, labels):
        """
        Args:
            x: feature matrix with shape (batch_size, feat_dim).
            labels: ground truth labels with shape (batch_size).
        """
        batch_size = x.size(0)
        # x = F.sigmoid(x)
        distmat = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(batch_size, self.cls_num) + \
                  torch.pow(self.centers, 2).sum(dim=1, keepdim=True).expand(self.cls_num, batch_size).t()
        distmat.addmm_(1, -2, x, self.centers.t())

        classes = torch.arange(self.cls_num).long()
        if self.use_gpu: classes = classes.cuda()
        labels = labels.unsqueeze(1).expand(batch_size, self.cls_num)
        mask = labels.eq(classes.expand(batch_size, self.cls_num))

        dist = distmat * mask.float()
        loss = dist.clamp(min=1e-12, max=1e+12).sum() / batch_size
        center = self.centers.detach().cpu().numpy()
        return loss,center